# !/usr/bin/env python
# -- coding: utf-8 --
# @Author zengxiaohui
# Datatime:7/13/2021 1:04 PM
# @File:大图切小图
import argparse
import os
import shutil

import cv2
from imutils import paths
import yaml
from tqdm import tqdm

from dataset.constants import LR_DATATYPE, TB_DATATYPE
from python_developer_tools.files.common import mkdir, get_filename_suf_pix
from python_developer_tools.python.multiprocessing_utils import parmap

splicing_col = [0.2575, 0.5019, 0.7450]  # 拼接位置的比例
splicing_row = [0.3372, 0.6489]  # 拼接位置在横向上的比例
lr_tb_thresh = 400

def cut_img(data_img_path):
    _, _, suffix, filename = get_filename_suf_pix(data_img_path)
    rgbimg = cv2.imread(data_img_path)
    height, width, _ = rgbimg.shape
    splicing_cols = [width * i for i in splicing_col]
    splicing_rows = [height * i for i in splicing_row]
    tb_pad_thresh = width / (24 * 2)  # 24列电池片
    lr_pad_thresh = height / (6 * 4)  # 6行电池片

    for i in range(len(splicing_cols) + 1):
        if i == 0:
            startx = 0
        else:
            startx = int(splicing_cols[i - 1] - tb_pad_thresh)

        if i >= len(splicing_cols):
            endx = width - 1
        else:
            endx = int(splicing_cols[i] + tb_pad_thresh)

        cv2.imwrite(os.path.join(tb_save_dir, '{}_{}_t{}'.format(filename, i, suffix)),
                    rgbimg[:lr_tb_thresh, startx:endx, :])
        cv2.imwrite(os.path.join(tb_save_dir, '{}_{}_b{}'.format(filename, i, suffix)),
                    rgbimg[height - lr_tb_thresh:height, startx:endx, :])

    for i in range(len(splicing_rows) + 1):
        if i == 0:
            starty = 0
        else:
            starty = int(splicing_rows[i - 1] - lr_pad_thresh)

        if i >= len(splicing_rows):
            endy = height - 1
        else:
            endy = int(splicing_rows[i] + lr_pad_thresh)

        cv2.imwrite(os.path.join(lr_save_dir, '{}_{}_l{}'.format(filename, i, suffix)),
                    rgbimg[starty:endy, :lr_tb_thresh, :])
        cv2.imwrite(os.path.join(lr_save_dir, '{}_{}_r{}'.format(filename, i, suffix)),
                    rgbimg[starty:endy, width - lr_tb_thresh:width - 1, :])

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="获取杭州人工认为有缺陷大图")
    parser.add_argument('--data_dir',
                        default=r"/home/zengxh/medias/data/ext/creepageDistance/20210714/bigimg/",
                        help="没有分的文件夹")
    parser.add_argument('--save_dir',
                        default=r"/home/zengxh/medias/data/ext/creepageDistance/20210714/smallimg",
                        help="收集好的图片保存在哪里")
    opt = parser.parse_args()

    data_img_paths = list(paths.list_images(opt.data_dir))
    lr_save_dir = os.path.join(opt.save_dir, LR_DATATYPE,"org")
    mkdir(lr_save_dir)
    tb_save_dir = os.path.join(opt.save_dir, TB_DATATYPE,"org")
    mkdir(tb_save_dir)

    parmap(cut_img, data_img_paths, 16)

